[ET-VK][sdpa] Add SDPA operator perf benchmark binary (test_sdpa)#21062
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21062
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ⏳ No Failures, 2 PendingAs of commit 7aa5888 with merge base 37400d9 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
400c19e
into
gh/SS-JIA/575/base
Pull Request resolved: #21062 Adds a standalone custom-ops perf binary that builds a ComputeGraph containing only the SDPA op (`llama.custom_sdpa`, the causal LLM KV-cache decode/prefill path) and reports per-dispatch GPU timing for the three SDPA dispatches (QK GEMM, softmax, AV GEMM) in isolation, so shader variants can be A/B-benchmarked without a full model. `TestSDPA.cpp` registers a `test_etvk.test_sdpa` op wrapper that synthesizes `input_pos` from the tensor shapes and drives `llama.custom_sdpa`. Two env toggles: `SDPA_DECODE_ONLY=1` restricts the sweep to the decode (S==1) shape matrix, and `SDPA_NO_CHAIN=1` sets `op_invocations_per_execute=1` to remove inter-invocation pipelining from the timing. Follows the existing `test/custom_ops/` pattern (e.g. `test_mm`) and is registered via `define_custom_op_test_binary` in `targets.bzl`. ghstack-source-id: 405400496 @exported-using-ghexport Differential Revision: [D112906312](https://our.internmc.facebook.com/intern/diff/D112906312/)
2 participants
Stack from ghstack (oldest at bottom):
Adds a standalone custom-ops perf binary that builds a ComputeGraph containing only the SDPA op (
llama.custom_sdpa, the causal LLM KV-cache decode/prefill path) and reports per-dispatch GPU timing for the three SDPA dispatches (QK GEMM, softmax, AV GEMM) in isolation, so shader variants can be A/B-benchmarked without a full model.TestSDPA.cppregisters atest_etvk.test_sdpaop wrapper that synthesizesinput_posfrom the tensor shapes and drivesllama.custom_sdpa. Two env toggles:SDPA_DECODE_ONLY=1restricts the sweep to the decode (S==1) shape matrix, andSDPA_NO_CHAIN=1setsop_invocations_per_execute=1to remove inter-invocation pipelining from the timing. Follows the existingtest/custom_ops/pattern (e.g.test_mm) and is registered viadefine_custom_op_test_binaryintargets.bzl.Differential Revision: D112906312